package com.example.attempt.service.impl;

import com.alibaba.fastjson.JSON;
import com.example.attempt.service.AIChatService;
import com.example.attempt.utils.WenXinChatUtils;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedList;
import java.util.Map;

@Slf4j
@Service
public class AIChatServiceImpl implements AIChatService {

    @Resource
    private WenXinChatUtils wenxinChatUtils;

    //历史对话，需要按照user,assistant 使用队列方便控制上下文长度
    LinkedList<Map<String,String>> messages = new LinkedList<>();

    @Override
    public String streamOutputToTerminal(String question) {
        //将问题放在历史对话中
        wenxinChatUtils.recordChatHistory(messages,"user",question);
        StringBuilder answer = new StringBuilder();
        // 发起异步请求
        Response response = wenxinChatUtils.getERNIEBot40ChatStream(1,messages,true);
        InputStream inputStream = null;
        ResponseBody responseBody = null;
        // 以流的方式处理响应内容，输出到控制台
        byte[] buffer = new byte[2048];
        int bytesRead;
        try {
            responseBody = response.body();
            if (responseBody != null) {
                inputStream = responseBody.byteStream();
            }
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                // 在控制台输出每个数据块
                System.out.write(buffer, 0, bytesRead);
                //将结果汇总起来  这里new出来的字符串中是包含id、token消耗数量等信息的JSON字符串，需要做信息统计的话可以在这里进行拓展
                answer.append(new String(buffer, 0, bytesRead));
            }
        } catch (IOException e) {
            //如果出现了异常就应该将问题也从对话历史中删除
            wenxinChatUtils.removeMessage(messages);
            log.error("InputStream流式读取出错 => {}",e.getMessage());
            throw new RuntimeException(e);
        }finally {
            wenxinChatUtils.closeStream(response,responseBody,inputStream);
        }

        //将回复的内容添加到消息中
        StringBuilder assistantAnswer = new StringBuilder();
        String[] answerArray = answer.toString().split("data: ");
        for (int i=1;i<answerArray.length;++i) {
            answerArray[i] = answerArray[i].substring(0,answerArray[i].length() - 2);
            assistantAnswer.append(JSON.parseObject(answerArray[i]).get("result"));
        }
        wenxinChatUtils.recordChatHistory(messages,"assistant",assistantAnswer.toString());
        return assistantAnswer.toString();
    }
}
